# -*- coding: utf-8 -*-

import tensorflow as tf
import h5py


def main():
    weights_name = ["conv1_1/weights", "conv1_2/weights",
                    "conv2_1/weights", "conv2_2/weights",
                    "conv3_1/weights", "conv3_2/weights", "conv3_3/weights", "conv3_4/weights"]
    biases_name = ["conv1_1/biases", "conv1_2/biases",
                   "conv2_1/biases", "conv2_2/biases",
                   "conv3_1/biases", "conv3_2/biases", "conv3_3/biases"]
    check_point_file_path = "../model/back/part_b_multi_1000/vgg-21000"
    reader = tf.train.NewCheckpointReader(check_point_file_path)
    h5_file = h5py.File("vgg-21000.h5", "w")
    for weight_name in weights_name:
        weight = reader.get_tensor(weight_name)
        h5_file.create_dataset(weight_name, data=weight)
    for bias_name in biases_name:
        bias = reader.get_tensor(bias_name)
        h5_file.create_dataset(bias_name, data=bias)
    h5_file.close()


if __name__ == "__main__":
    main()
